#include "agrad.hxx"
#include "utils.hxx"
#include "cuda/cuda_utils.cuh"
#include "cuda/helpers.cuh"

torch::Tensor
gcnhelper::gcn_norm_csr(
  torch::Tensor row_offset, 
  torch::Tensor adj_list, 
  torch::Tensor odegree, 
  torch::Tensor weight_list)
{
  CHECK_CUDA(row_offset);
  CHECK_CUDA(adj_list);
  CHECK_CUDA(odegree);
  CHECK_CUDA(weight_list);
  
  int nvertex = row_offset.size(0);
  auto ret = torch::zeros_like(weight_list);

  auto row_offset_ptr = row_offset.data_ptr<int>();
  auto adj_list_ptr = adj_list.data_ptr<int>();
  auto odegree_ptr = odegree.data_ptr<int>();

  AT_DISPATCH_FLOATING_TYPES(
  weight_list.scalar_type(), "gcnhelper::gcn_norm_csr", [&]{

    auto wl_ptr = weight_list.data_ptr<scalar_t>();
    auto out_ptr = ret.data_ptr<scalar_t>();
    
    const int tx = THD;
    const int bx = CEIL(nvertex, THD);
    CUDA_LAUNCH_CHECK(
      _gcn_normalize_csr_kernel<<<bx, tx>>>
      (row_offset_ptr, adj_list_ptr, odegree_ptr, wl_ptr, out_ptr, nvertex)
    );
    
  });

  return ret;
}

torch::Tensor
gcnhelper::gcn_norm_coo(
  torch::Tensor edge_list, 
  torch::Tensor weight_list,
  const int64_t nvertex)
{
  CHECK_CUDA(edge_list);
  CHECK_CUDA(weight_list);
  CHECK_INPUT(edge_list.dim() == 2);
  CHECK_INPUT(edge_list.size(0) == 2);

  const int nedge = edge_list.size(1);
  
  auto ret = torch::zeros_like(weight_list);

  auto dst_ptr = edge_list.data_ptr<int>();
  auto src_ptr = dst_ptr + nedge;

  int* odeg;
  // int *odeg_h;
  cudaMalloc(&odeg, sizeof(int)*nvertex);
  cudaMemset(odeg, 0, sizeof(int)*nvertex);
  // odeg_h = (int*) malloc(sizeof(int)*nvertex);

  CUDA_LAUNCH_CHECK(
    _degree_gather<<<CEIL(nvertex, THD), THD>>>(src_ptr, odeg, nvertex, nedge)
  );

  AT_DISPATCH_FLOATING_TYPES(
    weight_list.scalar_type(), "gcnhelper::gcn_norm_coo", [&]{

      auto wl_ptr = weight_list.data_ptr<scalar_t>();
      auto out_ptr = ret.data_ptr<scalar_t>();

      const int tx = THD;
      const int bx = CEIL(nedge, THD);

      CUDA_LAUNCH_CHECK(
        _gcn_normalize_coo_kernel<<<bx,tx>>>
        (src_ptr, dst_ptr, odeg, wl_ptr, out_ptr, nedge);
      );

  });

  // cudaMemcpy(odeg_h, odeg, sizeof(int)*nvertex, cudaMemcpyDeviceToHost);
  // for (int i=0; i<nvertex; ++i )
  // {
  //   printf("%d ", odeg_h[i]);
  // }
  // printf("\n");

  cudaFree(odeg);
  // free(odeg_h);
  return ret;
}